#!/usr/bin/env python3
"""Interactive helpers for searching World Bank indicators with wbgapi.

Usage:
  # run and type a query when prompted
  python scripts/search_wb.py

  # or run with query argument
  python scripts/search_wb.py --query "gdp per capita"
"""
import argparse
import wbgapi as wb


def search(query, limit=10):
    # wb.series.list supports a q= parameter for searching
    results = list(wb.series.list(q=query))
    if not results:
        print('No results')
        return
    for i, r in enumerate(results[:limit]):
        print(f"{i+1}. {r['id']} - {r.get('name')}")


def info(indicator_id):
    info = wb.series.info(indicator_id)
    print(info)


def list_countries(limit=20):
    for i, c in enumerate(wb.economy.list()):
        print(f"{c['id']} - {c['value']}")
        if i+1 >= limit:
            break


def sample(indicator_id, countries=['CHN','USA'], start=2010, end=2022):
    df = wb.data.DataFrame(indicator_id, economy=countries, time=range(start, end+1))
    print(df.head())


def main():
    p = argparse.ArgumentParser()
    p.add_argument('--query', '-q', help='Search query for indicators')
    p.add_argument('--info', '-i', help='Show series info for given indicator id')
    p.add_argument('--countries', '-c', nargs='+', help='List of countries for sample (default CHN USA)')
    p.add_argument('--sample', '-s', help='Get sample for indicator id', metavar='IND')
    args = p.parse_args()

    if args.query:
        search(args.query)
        return
    if args.info:
        info(args.info)
        return
    if args.sample:
        countries = args.countries if args.countries else ['CHN','USA']
        sample(args.sample, countries)
        return

    # interactive prompt
    q = input('Query (e.g. "gdp per capita"): ').strip()
    if q:
        search(q)


if __name__ == '__main__':
    main()
